# -*- coding: utf-8 -*-

from hlm.node import *
from hlm.base import Pack

from ip.fabrics import *
from ip.alu import ChannelAlu
from ip.erunner import ChannelERunnerSoftmem
from ip.opdef import AluOpCode


class PackElement(Pack):
    def __init__(self):
        self.sb: bool = False
        self.msgin: bool = False
        self.ain: bool = False
        self.data: complex = 0j


class ChannelShadowBufferFabric(Channel):
    lanes_map = {
        'Bus': {
            cp.PORTS: [('dst_uri', -1, int), ('r_data', -1, complex),
                       ('sb_addr', -1, int), ('ain_addr', -1, int)],
            cp.GET: ['dst_uri', 'r_data', 'sb_addr', 'ain_addr']
        }
    }


class NodeSoftmem(Node):
    def initialize(self, params: ObjectDict) -> None:
        params.softmem_size: int
        params.num_of_main_memory: int
        params.num_of_alu: int
        params.multiplexer: int
        self.const.num_of_alu = params.num_of_alu
        # Create the data holders.
        self.reg.data = [PackElement() for _ in range(params.softmem_size)]
        # Create the ALU request holder.
        self.reg.alu_dst_pid = [0 for _ in range(params.num_of_alu)]
        self.reg.alu_sb = [False for _ in range(params.num_of_alu)]
        self.reg.alu_dst_uri = [0 for _ in range(params.num_of_alu)]
        self.reg.alu_addr0 = [0 for _ in range(params.num_of_alu)]
        self.reg.alu_addr0_wait = [False for _ in range(params.num_of_alu)]
        self.reg.alu_addr0_ain = [False for _ in range(params.num_of_alu)]
        self.reg.alu_addr1 = [0 for _ in range(params.num_of_alu)]
        self.reg.alu_addr1_wait = [False for _ in range(params.num_of_alu)]
        self.reg.alu_addr1_ain = [False for _ in range(params.num_of_alu)]
        self.reg.alu_op_code = [AluOpCode.NOP for _ in range(params.num_of_alu)]
        # Create the ALU working states.
        self.reg.alu_has_work = [False for _ in range(params.num_of_alu)]
        self.reg.alu_working = [False for _ in range(params.num_of_alu)]
        # Save the bandwidth multiplexer.
        self.const.multiplexer = params.multiplexer
        # Create ports.
        self.port.memory_receivers = [self.port_channel_receiver(ChannelLoadFabric)
                                      for _ in range(params.num_of_main_memory)]
        self.port.memory_senders = [self.port_channel_sender(ChannelStoreFabric)
                                    for _ in range(params.num_of_main_memory)]
        # Prepare the M-Runner channel.
        self.port.mrunner_receiver = self.port_channel_receiver(ChannelLoadFabric)
        self.port.mrunner_sender = self.port_channel_sender(ChannelLoadFabric)
        # Prepare the E-Runner channel.
        self.port.alu_senders = [self.port_channel_sender(ChannelAlu)
                                 for _ in range(params.num_of_alu)]
        self.port.erunner_receivers = [self.port_channel_receiver(ChannelERunnerSoftmem)
                                       for _ in range(params.num_of_alu)]
        # Prepare the shadow buffer.
        self.port.shadow_buffer = self.port_channel_receiver(ChannelShadowBufferFabric)
        self.port.sb_save = self.port_channel_sender(ChannelShadowBufferFabric)
        ChannelShadowBufferFabric.connect(self.port.sb_save,
                                          self.port.shadow_buffer)
        # Shutdown flag.
        self.port.halt_alu = [self.port_out(bool)
                              for _ in range(params.num_of_alu)]

    tasks = ['task_main_memory_load',
             'task_shadow_buffer_load',
             'task_mrunner_load',
             'task_erunner_load',
             'task_alu_load',
             'task_deploy_request']

    def task_main_memory_load(self):
        for _ in range(self.const.multiplexer):
            # Process the load request, check all the memory channel.
            for memory_channel in self.port.memory_receivers:
                # Check the load request is empty or not.
                if memory_channel.Bus.is_any_empty():
                    continue
                # Check the ALU in request.
                self.reg.w_ain_addr = memory_channel.Bus.ain_addr.head()
                if self.reg.w_ain_addr != -1 and not self.reg.data[self.reg.w_ain_addr].ain:
                    continue
                # Check the SB in request.
                self.reg.w_sb_addr = memory_channel.Bus.sb_addr.head()
                if self.reg.w_sb_addr != -1 and not self.reg.data[self.reg.w_sb_addr].sb:
                    continue
                # All the condition is satisfied, set the request.
                _, self.reg.w_addr, self.reg.w_data, _, _, self.reg.w_clr = memory_channel.Bus.get()
                # Check the clear flag.
                if self.reg.w_clr:
                    self.reg.data[self.reg.w_addr].sb = False
                    self.reg.data[self.reg.w_addr].msgin = False
                    self.reg.data[self.reg.w_addr].ain = False
                    self.reg.data[self.reg.w_addr].data = 0j
                else:
                    self.reg.data[self.reg.w_addr].sb = False
                    self.reg.data[self.reg.w_addr].msgin = True
                    self.reg.data[self.reg.w_addr].ain = False
                    self.reg.data[self.reg.w_addr].data = self.reg.w_data

    def task_shadow_buffer_load(self):
        if self.port.shadow_buffer.Bus.is_any_empty():
            return
        # Check the ALU in request.
        self.reg.sb_ain_addr = self.port.shadow_buffer.Bus.ain_addr.head()
        if self.reg.sb_ain_addr != -1 and not self.reg.data[self.reg.sb_ain_addr].ain:
            return
        # Check the SB in request.
        self.reg.sb_sb_addr = self.port.shadow_buffer.Bus.sb_addr.head()
        if self.reg.sb_sb_addr != -1 and not self.reg.data[self.reg.sb_sb_addr].sb:
            return
        # Now the data could be set to soft-mem.
        self.reg.sb_w_addr, self.reg.sb_w_data, _, _ = self.port.shadow_buffer.Bus.get()
        self.reg.data[self.reg.sb_w_addr].sb = False
        self.reg.data[self.reg.sb_w_addr].msgin = True
        self.reg.data[self.reg.sb_w_addr].ain = False
        self.reg.data[self.reg.sb_w_addr].data = self.reg.sb_w_data

    def task_mrunner_load(self):
        # Process the input data from mrunner.
        if self.port.mrunner_receiver.Bus.is_any_empty():
            return
        # Check the ALU in request.
        self.reg.m_ain_addr = self.port.mrunner_receiver.Bus.ain_addr.head()
        if self.reg.m_ain_addr != -1 and not self.reg.data[self.reg.m_ain_addr].ain:
            return
        # Check the SB in request.
        self.reg.m_sb_addr = self.port.mrunner_receiver.Bus.sb_addr.head()
        if self.reg.m_sb_addr != -1 and not self.reg.data[self.reg.m_sb_addr].sb:
            # Put the data to shadow buffer.
            if not self.port.shadow_buffer.Bus.is_any_full():
                # Save the data to shadow buffer.
                _, self.reg.m_addr, self.reg.m_data, _, _, _ = \
                    self.port.mrunner_receiver.Bus.get()
                self.port.sb_save.Bus.put(dst_uri=self.reg.m_addr,
                                          r_data=self.reg.m_data,
                                          sb_addr=self.reg.m_sb_addr,
                                          ain_addr=self.reg.m_ain_addr)
        else:
            # Update the data.
            _, self.reg.m_addr, self.reg.m_data, _, _, _ = \
                self.port.mrunner_receiver.Bus.get()
            # Update the data.
            self.reg.data[self.reg.m_addr].sb = False
            self.reg.data[self.reg.m_addr].msgin = True
            self.reg.data[self.reg.m_addr].ain = False
            self.reg.data[self.reg.m_addr].data = self.reg.m_data

    def task_erunner_load(self):
        # Load the request from ERunner.
        for ii in range(self.const.num_of_alu):
            # Check whether the working request is empty.
            # If so, now we need to check whether we have request coming from
            # E-Runner.
            if not self.reg.alu_has_work[ii] and \
                    not self.port.erunner_receivers[ii].Read.is_any_empty():
                # Save the new ALU working request.
                self.reg.alu_dst_pid[ii], self.reg.alu_sb[ii], \
                self.reg.alu_dst_uri[ii],  self.reg.alu_addr0[ii], \
                self.reg.alu_addr0_wait[ii], self.reg.alu_addr0_ain[ii], \
                self.reg.alu_addr1[ii], self.reg.alu_addr1_wait[ii], \
                self.reg.alu_addr1_ain[ii], self.reg.alu_op_code[ii] = \
                    self.port.erunner_receivers[ii].Read.get()
                # Set the has work flag.
                self.reg.alu_has_work[ii] = True

    def task_alu_load(self):
        # Load the request from ALU.
        for ii in range(self.const.num_of_alu):
            # Check whether there is a writing request from ALU.
            if self.port.alu_senders[ii].Dst.is_any_empty():
                continue
            # Fetch the loading request from ALU.
            self.reg.alu_dst_addr, self.reg.alu_dst_data = \
                self.port.alu_senders[ii].Dst.get()
            # Save the ALU data.
            self.reg.data[self.reg.alu_dst_addr].sb = False
            self.reg.data[self.reg.alu_dst_addr].msgin = False
            self.reg.data[self.reg.alu_dst_addr].ain = True
            self.reg.data[self.reg.alu_dst_addr].data = self.reg.alu_dst_data
            # Clear the waiting and has work flag.
            self.reg.alu_has_work[ii] = False
            self.reg.alu_working[ii] = False

    def task_deploy_request(self):
        # Check whether we could deploy the data.
        for ii in range(self.const.num_of_alu):
            # Check whether ALU got its work.
            if not self.reg.alu_has_work[ii]:
                continue
            # Check the op code.
            # Three kinds of request is detected:
            #  - ALU op.
            #  - Msg op.
            #  - Halt op.
            if self.reg.alu_op_code[ii] == AluOpCode.HALT:
                # Set the shutdown signal.
                self.port.halt_alu[ii].put(True)
                # Clear the has work flag.
                self.reg.alu_has_work[ii] = False
                continue
            if self.reg.alu_op_code[ii] < AluOpCode.PUT:
                # Check whether the ALU sender is ready.
                if self.reg.alu_working[ii] or \
                        self.port.alu_senders[ii].Src.is_any_full():
                    continue
                # Only when the following condition is full-filled:
                # - If the data need to be waited, then wait it.
                # - If the data is ready, then check next.
                # Check wait condition.
                if self.reg.alu_addr0_wait[ii] and not self.reg.data[self.reg.alu_addr0[ii]].msgin:
                    continue
                if self.reg.alu_addr1_wait[ii] and not self.reg.data[self.reg.alu_addr1[ii]].msgin:
                    continue
                # Check ALU in condition.
                if self.reg.alu_addr0_ain[ii] and not self.reg.data[self.reg.alu_addr0[ii]].ain:
                    continue
                if self.reg.alu_addr1_ain[ii] and not self.reg.data[self.reg.alu_addr1[ii]].ain:
                    continue
                # Release the data message in flag.
                self.reg.data[self.reg.alu_addr0[ii]].msgin = False
                self.reg.data[self.reg.alu_addr1[ii]].msgin = False
                # Mark the ALU is working.
                self.reg.alu_working[ii] = True
                # Deploy the mission to ALU.
                self.port.alu_senders[ii].Src.put(
                    addr=self.reg.alu_dst_uri[ii],
                    data0=self.reg.data[self.reg.alu_addr0[ii]].data,
                    data1=self.reg.data[self.reg.alu_addr1[ii]].data,
                    op=self.reg.alu_op_code[ii])
                continue
            if AluOpCode.PUT <= self.reg.alu_op_code[ii] < AluOpCode.HALT:
                # The message OP might goes directly to memory or hold and
                # wait here.
                if self.reg.alu_op_code[ii] == AluOpCode.WAIT:
                    # Hold and wait, only check the msgin bit, but not do
                    # anything.
                    if self.reg.data[self.reg.alu_addr0[ii]].msgin:
                        # Condition satisfied, clear the request
                        self.reg.alu_has_work[ii] = False
                elif self.reg.alu_op_code[ii] == AluOpCode.PUT:
                    # Pass to MRunner or main memory, based on the destination pid.
                    if self.reg.alu_dst_pid[ii] == -1:
                        # Pass to main memory.
                        # Check wait condition.
                        if self.reg.alu_addr0_wait[ii] and not self.reg.data[self.reg.alu_addr0_wait[ii]].msgin:
                            continue
                        # Will pass to main memory write, finding a free memory controller.
                        for memory_channel in self.port.memory_senders:
                            if memory_channel.Bus.is_any_full():
                                continue
                            # Send to the current channel.
                            memory_channel.Bus.put(addr=self.reg.alu_dst_uri[ii],
                                                   data=self.reg.data[self.reg.alu_addr0[ii]].data)
                            # Set the shadow buffer bit.
                            self.reg.data[self.reg.alu_addr0[ii]].sb = True
                            self.reg.data[self.reg.alu_addr0[ii]].msgin = False
                            # Clear the request.
                            self.reg.alu_has_work[ii] = False
                    else:
                        # Pass to M-Runner.
                        if self.port.mrunner_sender.Bus.is_any_full():
                            continue
                        # Deploy the data.
                        self.port.mrunner_sender.Bus.put(
                            pid=self.reg.alu_dst_pid[ii],
                            addr=self.reg.alu_dst_uri[ii],
                            data=self.reg.data[self.reg.alu_addr0[ii]].data,
                            sb_addr=self.reg.alu_dst_uri[ii] if self.reg.alu_sb[ii] else -1,
                            ain_addr=-1,
                            clr=False)
                        # Update the shadow buffer bit.
                        self.reg.data[self.reg.alu_addr0[ii]].sb = True
                        self.reg.data[self.reg.alu_addr0[ii]].msgin = False
                        # Clear the working request.
                        self.reg.alu_has_work[ii] = False
